import requests
import time
import jwt
import os

class GlmApi:

    def __init__(self, model='glm-4', system_prompt=None):
        self.model_name = model
        self.system_prompt = system_prompt if system_prompt else '你是一个乐于解答各种问题的助手，你的任务是为用户提供专业、准确、有见地的建议。'
        self.model_url = 'https://open.bigmodel.cn/api/paas/v4/chat/completions'
        self.api_key = os.environ['ZHIPUAI_API_KEY']

    def request(self, text, debug=False):
        if self.api_key is None:
            return "Api key is required.", False
        try:
            headers = self.generate_headers()
            payload = self.generate_payload(text)
            response = requests.post(self.model_url, headers=headers, json=payload)
            response.raise_for_status()
            ans = response.json()
            if debug:
                print(ans)
            result = ans["choices"][0]["message"]["content"]
            return result, True
        except requests.exceptions.RequestException as e:
            raise Exception(f"请求异常：{e}")
        except requests.exceptions.Timeout as e:
            raise Exception(f"请求超时：{e}")
        except Exception as e:
            raise Exception(f"发生了未知错误：{e}")
        return "", False

    def generate_token(self, apikey: str, exp_seconds: int = 3600*24):
        try:
            id, secret = apikey.split(".")
        except Exception as e:
            raise Exception("invalid apikey", e)

        payload = {
            "api_key": id,
            "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
            "timestamp": int(round(time.time() * 1000)),
        }

        return jwt.encode(
            payload,
            secret,
            algorithm="HS256",
            headers={"alg": "HS256", "sign_type": "SIGN"},
        )

    def generate_headers(self):
        token = self.generate_token(self.api_key)
        return {
            'Authorization': f'Bearer {token}',
            'Content-Type': 'application/json'
        }

    def generate_payload(self, text):
        return {
            'model': self.model_name,
            'messages': [
                {
                    'role': 'system',
                    'content': self.system_prompt
                },
                {
                    'role': 'user',
                    'content': text
                }
            ]
        }
    

class NovelCharecter(GlmApi):
    """小说角色"""

    def __init__(self, character, opponent):
        super().__init__()
        self.model_name = 'charglm-3'
        self.character = character
        self.opponent = opponent
        self.messages = []
    
    def generate_payload(self, text):
        message = {
            'role': 'user',
            'content': text
        }
        self.messages.append(message)
        return {
            'model': self.model_name,
            'meta': {
                "user_info": self.character['description'],
                "user_name": self.character['name'],
                "bot_info": self.opponent['description'],
                "bot_name": self.opponent['name']
            },
            'messages': self.messages
        }


def test_api():
    glm = GlmApi()
    result, succ = glm.request("请问地球到月球距离有多远", debug=True)
    print(result)


if __name__ == '__main__':
    test_api()